import tensorflow as tf
from tensorflow import keras
#报错不影响使用
import numpy as np
import matplotlib.pyplot as plt

# 加载 MNIST 数据集
mnist = keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# 归一化处理（像素值 0-255 映射到 0-1）
x_train, x_test = x_train / 255.0, x_test / 255.0

# 构建 CNN 模型
model = keras.Sequential([
    keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(28,28,1)),
    keras.layers.MaxPooling2D((2,2)),
    keras.layers.Conv2D(64, (3,3), activation='relu'),
    keras.layers.MaxPooling2D((2,2)),
    keras.layers.Flatten(),
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dense(10, activation='softmax')
])

# 编译模型
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 训练模型
model.fit(x_train.reshape(-1, 28, 28, 1), y_train, epochs=5, validation_split=0.1)

# 在测试集上评估
test_loss, test_acc = model.evaluate(x_test.reshape(-1, 28, 28, 1), y_test)
print(f"Test Accuracy: {test_acc:.4f}")

# 预测一张测试图片
index = np.random.randint(0, len(x_test))
plt.imshow(x_test[index], cmap='gray')
plt.show()
pred = model.predict(x_test[index].reshape(1, 28, 28, 1))
print(f"Predicted Label: {np.argmax(pred)}")
